import pandas as pd, numpy as np
from sklearn.feature_extraction.text import CountVectorizer, TfidfVectorizer
from sklearn import svm
import pickle

data = pickle.load(open('./train_tf_idf.pkl','rb'))
train_word,train_classify,test_word,test_id,trn_term_doc,test_term_doc = data[0],data[1],data[2],data[3],data[4],data[5]

fid0=open('svm_result.csv','w')
train_classify = np.array(train_classify)
y=(train_classify-1).astype(int)
lin_clf = svm.LinearSVC()
lin_clf.fit(trn_term_doc,y)
preds = lin_clf.predict(test_term_doc)
i=0
fid0.write("id,class"+"\n")
for item in preds:
    fid0.write(str(i)+","+str(item+1)+"\n")
    i=i+1
fid0.close()